package com.hundun.vision.biz.core;

import org.ta4j.core.BarSeries;
import org.ta4j.core.BaseTradingRecord;
import org.ta4j.core.Trade.TradeType;
import org.ta4j.core.TradingRecord;
import org.ta4j.core.analysis.cost.CostModel;
import org.ta4j.core.num.DecimalNum;

import java.math.BigDecimal;

import static java.math.RoundingMode.DOWN;
import static org.ta4j.core.Trade.TradeType.BUY;
import static org.ta4j.core.Trade.TradeType.SELL;

/**
 * @author ：RuoChen
 * @date ：15:45 2024/12/6
 * @description：合约回测管理器
 */

public class ContractBarSeriesManager {
    /**
     * K线数据系列
     */
    private BarSeries series;
    /**
     * 合约策略
     */
    private ContractStrategy strategy;

    /**
     * 交易成本模型
     */
    private CostModel transactionCostModel;

    /**
     * 交易成本模型
     */
    private CostModel holdingCostModel;
    /**
     * 交易记录
     */
    private TradingRecord tradingRecord;
    /**
     * 当前账户余额
     */
    private BigDecimal balance;
    /**
     * 当前仓位
     */
    private BigDecimal position;
    /**
     * 杠杆倍数
     */
    private BigDecimal leverage;

    public ContractBarSeriesManager(BarSeries series, ContractStrategy strategy, BigDecimal leverage, CostModel transactionCostModel, CostModel holdingCostModel) {
        this.series = series;
        this.strategy = strategy;
        this.transactionCostModel = transactionCostModel;
        this.holdingCostModel = holdingCostModel;
        this.leverage = leverage;
        this.balance = BigDecimal.valueOf(100);
        this.position = BigDecimal.ZERO;
    }

    public TradingRecord run() {
        for (int index = 0; index < series.getBarCount(); index++) {
            BigDecimal price = series.getBar(index).getClosePrice().bigDecimalValue();
            //多单/空单建仓
            ContractRule entryRule = strategy.getEntryRule(index, tradingRecord);
            if (entryRule != null) {
                boolean isBuy = entryRule.isBuy();
                if (isBuy && position.compareTo(BigDecimal.ZERO) == 0) {
                    if (tradingRecord == null) {
                        this.tradingRecord = new BaseTradingRecord(BUY, transactionCostModel, holdingCostModel);
                    }
                    // 多单建仓
                    BigDecimal quantity = balance.divide(price, 6, DOWN).multiply(leverage).setScale(6, DOWN);
                    System.out.println("Buy entry at index:" + index + " price:" + price + " quantity:" + quantity);
                    // 记录买入操作
                    tradingRecord.operate(index, DecimalNum.valueOf(price), DecimalNum.valueOf(quantity));
                    position = quantity;
                    balance = balance.subtract(price.multiply(quantity).divide(leverage, 6, DOWN).setScale(6, DOWN));
                } else if (!isBuy && position.compareTo(BigDecimal.ZERO) == 0) {
                    if (tradingRecord == null) {
                        this.tradingRecord = new BaseTradingRecord(SELL, transactionCostModel, holdingCostModel);
                    }
                    // 空单建仓
                    BigDecimal quantity = balance.divide(price, 6, DOWN).multiply(leverage).setScale(6, DOWN);
                    System.out.println("Sell entry at index:" + index + " price:" + price + " quantity: " + quantity);
                    // 记录卖出操作
                    tradingRecord.operate(index, DecimalNum.valueOf(price), DecimalNum.valueOf(quantity));
                    position = quantity.negate();
                    balance = balance.subtract(price.multiply(quantity).divide(leverage, 6, DOWN)).setScale(6, DOWN);
                }
            } else {
                //出场信号：止盈/止损
                ContractRule exitRule = strategy.getExitRule(index, tradingRecord);
                if (exitRule != null) {
                    boolean isBuy = exitRule.isBuy();
                    if (isBuy && position.compareTo(BigDecimal.ZERO) < 0) {
                        // 空单止盈/止损
                        System.out.println("Sell exit at index:" + index + " price:" + price + " quantity:" + position);
                        // 记录买入操作
                        balance = getBalance(balance, SELL, tradingRecord.getLastTrade().getPricePerAsset().bigDecimalValue(), price, position.abs(), leverage);
                        tradingRecord.operate(index, DecimalNum.valueOf(price), DecimalNum.valueOf(position.abs()));
                        position = BigDecimal.ZERO;
                    } else if (!isBuy && position.compareTo(BigDecimal.ZERO) > 0) {
                        // 多单止盈/止损
                        System.out.println("Buy exit at index:" + index + " price:" + price + " quantity:" + position);
                        // 记录卖出操作
                        balance = getBalance(balance, BUY, tradingRecord.getLastTrade().getPricePerAsset().bigDecimalValue(), price, position, leverage);
                        tradingRecord.operate(index, DecimalNum.valueOf(price), DecimalNum.valueOf(position));
                        position = BigDecimal.ZERO;
                    }
                }
            }
        }
        return tradingRecord;
    }

    /**
     * 获取账户余额
     *
     * @param tradeType
     * @param lastPrice
     * @param currentPrice
     * @param quantity
     * @return
     */
    private BigDecimal getBalance(BigDecimal balance, TradeType tradeType, BigDecimal lastPrice,
                                  BigDecimal currentPrice, BigDecimal quantity, BigDecimal leverage) {
        balance = balance.add(lastPrice.multiply(quantity).divide(leverage));
        BigDecimal profitOrLoss = lastPrice.subtract(currentPrice).multiply(quantity);
        if (BUY.equals(tradeType)) {
            //多单止损后余额
            if (profitOrLoss.compareTo(BigDecimal.ZERO) >= 0) {
                balance = balance.subtract(profitOrLoss);
            } else {
                //多单止盈后余额
                balance = balance.add(profitOrLoss.negate());
            }
        } else {
            //空单止损后余额
            if (profitOrLoss.compareTo(BigDecimal.ZERO) <= 0) {
                balance = balance.subtract(profitOrLoss.negate());
            } else {
                //空单止盈后余额
                balance = balance.add(profitOrLoss);
            }
        }
        balance.setScale(6, DOWN);
        return balance;
    }
}